"""
RFdiffusion图操作算子包
"""

from ops.src.edge_softmax.edge_softmax import edge_softmax
from ops.src.e_dot_v.e_dot_v import e_dot_v
from ops.src.copy_e_sum.copy_e_sum import copy_e_sum
from ops.src.gspmm.gspmm import gspmm
from ops.src.gsddmm.gsddmm import gsddmm
from ops.src.graph.graph import graph, Graph

# 兼容DGL的命名空间
class OpsNamespace:
    """模拟DGL的ops命名空间"""
    @staticmethod
    def edge_softmax(graph, edge_weights):
        return edge_softmax(graph, edge_weights)
    
    @staticmethod
    def e_dot_v(graph, edge_features, node_features):
        return e_dot_v(graph, edge_features, node_features)
    
    @staticmethod
    def copy_e_sum(graph, edge_features):
        return copy_e_sum(graph, edge_features)
    
    @staticmethod
    def gspmm(graph, op, reduce_op, lhs_data, rhs_data):
        return gspmm(graph, op, reduce_op, lhs_data, rhs_data)
    
    @staticmethod
    def gsddmm(graph, op, lhs_data, rhs_data, lhs_target='u', rhs_target='v'):
        return gsddmm(graph, op, lhs_data, rhs_data, lhs_target, rhs_target)

# 导出ops命名空间
ops = OpsNamespace() 